#include "TFile.h"
#include "TString.h"
#include "TTree.h"
#include "TROOT.h"

#include "InputFileReader.hh"
#include "OutputFileWriter.hh"
#include "PMT.hh"
#include "PMTContainer.hh"

#include <string>
#include <cmath>
#include <filesystem>

OutputFileWriter::OutputFileWriter(YAML::Node configer)
{
    //m_VSumQ.clear();
    //m_VPMTLineShape.clear();
    TString OutName = configer["OutputName"].as<std::string>();
    m_DrawWhileDigi = configer["DrawWhileDigi"].as<bool>();
    m_OutFile = new TFile(OutName, "recreate");
    m_OutTree = new TTree("DigiInfo", "");

    m_OutTree->Branch("PGTypeFlag",     &m_PGTypeFlag,      "PGTypeFlag/I");
    m_OutTree->Branch("TruthEnergy",    &m_TruthEnergy,     "TruthEnergy/D");
    m_OutTree->Branch("TruthPos",       &m_TruthPos,        "TruthPos[3]/D");
    m_OutTree->Branch("TruthMomDir",    &m_TruthMomDir,     "TruthMomDir[3]/D");
    m_OutTree->Branch("NumOfPMT",       &m_NumOfPMT,        "NumOfPMT/I");
    m_NumOfPMT = PMTContainer::GetInstance()->NumOfPMT();
    if(m_NumOfPMT > MAXPMT)
    {
        std::cerr << "Error!!! Number of PMT larger than MAXPMT in OutputFileWriter.hh. Please enlarge it!!!" << std::endl;
        exit(1);
    }
    for(int i = 0; i < m_NumOfPMT; ++i)
    {
        //m_VSumQ.push_back(0);
        //m_VPMTLineShape.push_back(new TGraph());
        m_VSumQ[i] = 0;
        m_VPMTLineShape[i] = new TGraph();
        m_OutTree->Branch(TString::Format("SumQ_%d", i),    &(m_VSumQ[i]),    TString::Format("SumQ_%d/D", i));
        m_OutTree->Branch(TString::Format("PMTLineShape_%d", i), &(m_VPMTLineShape[i]), 64000, 0);
        //TBranch* br_Graph = m_OutTree->Branch(TString::Format("PMTLineShape_%d", i), &m_VPMTLineShape[i], 32000, 0);
        //br_Graph->SetAutoDelete(true);
    }

    if(m_DrawWhileDigi)
    {
        Int_t row = configer["DrawOption"]["Row"].as<int>();
        Int_t line = ceil(m_NumOfPMT / double(row));
        gROOT->ProcessLine(".L lhcbStyle.C");
        m_Can = new TCanvas("can", "", 400*row, 300*line);
        m_Can->Divide(row, line);

        m_iPlot = 0;
        m_NumOfPlots = 0;
        if(configer["DrawOption"]["NumOfPlots"].IsDefined())
        {
            m_NumOfPlots = configer["DrawOption"]["NumOfPlots"].as<int>();
        }
        m_DrawPath = configer["DrawOption"]["Path"].as<std::string>();
        if(!std::filesystem::exists(m_DrawPath))
        {
            std::filesystem::create_directory(m_DrawPath);
        }
    }
}

OutputFileWriter::~OutputFileWriter()
{
    m_OutFile->Close();
}

void OutputFileWriter::SaveAll()
{
    m_OutFile->cd();
    m_OutTree->AutoSave();
}

void OutputFileWriter::WriteEvent(aSimEvent& aEvent)
{
    m_PGTypeFlag = aEvent.Type;
    m_TruthEnergy = aEvent.Energy;
    for(int i = 0; i < 3; ++i)
    {
        m_TruthPos[i] = aEvent.Pos[i];
        m_TruthMomDir[i] = aEvent.MomDir[i];
    }

    for (int i = 0; i < m_NumOfPMT; ++i)
    {
        const PMT* ptrPMT = PMTContainer::GetInstance()->GetPMT(i);
        const std::pair<Double_t, TGraph*> result = ptrPMT->Result();

        m_VSumQ[i] = result.first;
        m_VPMTLineShape[i] = result.second;
    }
    if(m_DrawWhileDigi && (m_NumOfPlots > 0) && (m_iPlot < m_NumOfPlots))
    {
        for (int i = 0; i < m_NumOfPMT; ++i)
        {
            m_Can->cd(i+1);
            m_VPMTLineShape[i]->SetLineWidth(1);
            m_VPMTLineShape[i]->GetXaxis()->SetTitle("Time (ns)");
            m_VPMTLineShape[i]->GetYaxis()->SetTitle("SigV (ADC)");
            m_VPMTLineShape[i]->Draw("AL");
        }
        m_Can->SaveAs(TString::Format("%s/Event_%lld.png", m_DrawPath.c_str(), m_OutTree->GetEntries()));
        ++m_iPlot;            
    }
    m_OutTree->Fill();
}